#!/usr/bin/env python3

# Copyright (c) 2018, Kontron Europe GmbH
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

from __future__ import print_function

# to disable warnings generated by building the font cache
import warnings
warnings.filterwarnings('ignore')

import argparse
import json
import numpy as np
import sys

from os import environ
if environ.get('DISPLAY') is None:
    import matplotlib as mpl
    mpl.use('Agg')
import matplotlib.pyplot as plt

def chunks(d, n):
    for i in range(0, len(d), n):
        yield d[i:i+n]

def reduce_sum_array(a, n):
    return [sum(l) for l in chunks(a, n)]

def add_subplot_data(plt, data, unit, width, reduce=1):
    offset = 0
    max_val = data['max']
    min_val = data['min']
    outliers = data['outliers']

    y = data['histogram']

    # reduce data to show in plot
    if reduce != 1:
        y = reduce_sum_array(y, reduce)

    if 'offset' in data:
        offset = data['offset']
        x = range(int(-offset/reduce), len(y) - int(offset/reduce))
    else:
        x = range(0, len(y))

    if reduce != 1:
        x = [e * reduce for e in x]

    plt.bar(x, y, width, edgecolor="none")
    plt.yscale('symlog')
    ax = plt.gca()
    #ax.set_xticks([0,100,200,300,400,500,600,700,800,900,1000])
    #ax.set_xticklabels(range(0, len(x) * reduce))
    t = 'min: %s $\mu$s\nmax: %s $\mu$s\noutliers: %s' % (
            min_val, max_val, outliers)
    text = 'min: %s %s\nmax: %s %s\noutliers: %s' % (
            min_val, unit,  max_val, unit, outliers)
    plt.text(1, 1, text, transform=ax.transAxes, verticalalignment='top',
            horizontalalignment='right')


def plot(rxtimestamps, txtimestamps, jitter, props):
    plt.figure(figsize=(10,8))
    plt.suptitle(props['plottitle'], fontsize=14, fontweight='bold')

    plt.subplot(2, 2, (1,2))
    plt.title('Scheduled times')
    plt.grid(True)
    plt.xlabel('Latency ($\mu$s)')
    plt.ylabel('Packets')
    add_subplot_data(plt, rxtimestamps, '$\mu$s', width=2, reduce=2)

    plt.subplot(2, 2, 3)
    plt.title('RT Application latency')
    plt.grid(True)
    plt.xlabel('Latency ($\mu$s)')
    plt.ylabel('Samples')
    add_subplot_data(plt, txtimestamps, '$\mu$s', width=0.8)

    plt.subplot(2, 2, 4)
    plt.title('TSN Network Jitter')
    plt.grid(True)
    plt.xlabel('Jitter (ns)')
    plt.ylabel('Packets')
    add_subplot_data(plt, jitter, 'ns', width=10, reduce=10)

    plt.subplots_adjust(top=0.90, hspace=0.4, bottom=0.15)
    counts = rxtimestamps['count']
    start = rxtimestamps['start-timestamp']
    end = rxtimestamps['end-timestamp']
    plt.gcf().text(0.01, 0.01, 'counts: %s\nstart: %s\nend: %s' % (counts, start, end))

def plot_output(outfile=None):
    if outfile:
        plt.savefig(outfile)
    else:
        try:
            plt.show()
        except KeyboardInterrupt:
            pass


def main(args=None):
    parser = argparse.ArgumentParser(
        description='reportgen')
    parser.add_argument('--title', dest='plottitle', type=str,
                        default='TSN latency and jitter report',
                        help='Set plot title.')
    parser.add_argument('infile', nargs='?', type=argparse.FileType('r'),
                        help='Input file (default is STDIN)', default=sys.stdin)
    parser.add_argument('outfile', nargs='?', type=str,
                        help='Output file (default is None, open X window)',
                        default=None)
    args = parser.parse_args(args)

    props = dict(plottitle=args.plottitle)

    j = None
    a = None
    b = None
    c = None
    try:
        for line in args.infile:
            line = line.strip()
            if not line:
                continue
            try:
                j = json.loads(line)
                if j['type'] == 'histogram-program-latency':
                    b = j['object']
                if j['type'] == 'histogram-scheduled-times':
                    a = j['object']
                if j['type'] == 'histogram-jitter':
                    c = j['object']
            except ValueError as e:
                print(e)
                pass
    except KeyboardInterrupt:
        import time
        time.sleep(0.5)
        for line in args.infile:
            line = line.strip()
            if not line:
                continue
            try:
                j = json.loads(line)
                if j['type'] == 'histogram-program-latency':
                    b = j['object']
                if j['type'] == 'histogram-scheduled-times':
                    a = j['object']
                if j['type'] == 'histogram-jitter':
                    c = j['object']
            except ValueError:
                pass

    plot(a, b, c, props)
    plot_output(args.outfile)


if __name__ == '__main__':
    main()
